import cv2
import numpy as np


# 定义一个函数，将YOLO的标签格式转换为左上角和右下角的坐标
def xywh2xyxy(x, w, h, img):
    label, x, y, w, h = x
    # 反归一化并得到左上和右下坐标
    x_t = x * img.shape[1]
    y_t = y * img.shape[0]
    w_t = w * img.shape[1]
    h_t = h * img.shape[0]
    top_left_x = int(x_t - w_t / 2)
    top_left_y = int(y_t - h_t / 2)
    bottom_right_x = int(x_t + w_t / 2)
    bottom_right_y = int(y_t + h_t / 2)
    return (top_left_x, top_left_y), (bottom_right_x, bottom_right_y)


# 读取图像和标签
def visualize_yolo_labels(image_path, label_path, classes):
    img = cv2.imread(image_path)
    height, width = img.shape[:2]

    with open(label_path, 'r') as f:
        labels = f.read().splitlines()

    for label in labels:
        class_id, x_center, y_center, width, height = map(float, label.split())
        class_id = int(class_id)
        # 将中心点坐标和宽高转换为左上角和右下角坐标
        (top_left_x, top_left_y), (bottom_right_x, bottom_right_y) = xywh2xyxy(
            (class_id, x_center, y_center, width, height), width, height, img)

        # 绘制矩形框和标签
        cv2.rectangle(img, (top_left_x, top_left_y), (bottom_right_x, bottom_right_y), (0, 255, 0), 2)
        cv2.putText(img, classes[class_id], (top_left_x, top_left_y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255),
                    2)

    # 显示图像
    cv2.imshow('Image with YOLO labels', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


# 类别列表
classes = ['dog']  # 请替换为您的类别名称

# 图像路径和标签路径
image_path = 'D:/Smart Video Editing/video/second/dataset/images/image_1.png'  # 替换为您的图像路径
label_path = 'D:/Smart Video Editing/video/second/dataset/labels/image_1.txt'  # 替换为您的标签路径

# 可视化标签
visualize_yolo_labels(image_path, label_path, classes)

